/*
 * Copyright 2012 David Crosson
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package fr.janalyse.jmx

import scala.util.{Properties=>SP}
import java.io.File.{separator=>FS, pathSeparator=>PS}
import scala.reflect.BeanProperty
import java.lang.management.ManagementFactory
import javax.management.ObjectName
import javax.management.ObjectInstance
import javax.management.ObjectName
import javax.management.remote.rmi.RMIConnection
import javax.management.remote.rmi.RMIServer
import javax.management.remote.rmi.RMIServerImpl_Stub
import javax.management.openmbean.{CompositeData, CompositeDataSupport}
import javax.management.openmbean.{TabularData,TabularDataSupport}
import scala.collection.JavaConversions._
import java.rmi.MarshalledObject
import java.rmi.UnmarshalException
import javax.management.openmbean.CompositeData
import scala.collection.JavaConversions._


trait JMXAutoClose {
  def usingJMX[T <: { def close() }, R](resource: T)(block: T => R) = {
    try block(resource)
    finally resource.close
  }
}


case class CompositeDataWrapper(cd:CompositeData) {
  // TODO : BAD because it assumes that all values are of the same type...
  def toMap[RT]() = {
    val keys = cd.getCompositeType().keySet()
    val tuples = keys.toList map {k => (k -> cd.get(k).asInstanceOf[RT]) }
    tuples.toMap
  }
}

case class CompositeDataSupportWrapper(cd:CompositeDataSupport) {
  // TODO : BAD because it assumes that all values are of the same type...
  def toMap[RT]() = {
    val keys = cd.getCompositeType().keySet()
    val tuples = keys.toList map {k => (k -> cd.get(k).asInstanceOf[RT]) }
    tuples.toMap
  }
}

case class TabularDataWrapper(tabularData:TabularData) {
  def toMap[X]() = {
      val indexNamesInCell = tabularData.getTabularType().getIndexNames()
	  val tuples = tabularData.values() collect {
		case (k:java.util.List[String], v:CompositeData) => {
			val cellContent = CompositeDataWrapper(v).toMap[X]
			val key    = indexNamesInCell map {cellContent.get(_).get} mkString("-")
			val submap = cellContent filterNot {case (k,_) => indexNamesInCell contains k}
			
			key -> submap
		}
	  }
      tuples.toMap
  }
}


case class TabularDataSupportWrapper(tabularData:TabularDataSupport) {
  def toMap[X]() = {
      val indexNamesInCell = tabularData.getTabularType().getIndexNames()
	  val tuples = tabularData collect {
		case (k:java.util.List[String], v:CompositeDataSupport) => {
			val cellContent = CompositeDataSupportWrapper(v).toMap[X]
			val key    = indexNamesInCell map {cellContent.get(_).get} mkString("-")
			val submap = cellContent filterNot {case (k,_) => indexNamesInCell contains k}
			
			key -> submap
		}
	  }
      tuples.toMap
  }
}


trait JMXImplicits {
  implicit def string2objectName(name:String):ObjectName = new ObjectName(name)
  implicit def objectInstance2ObjectName(that:ObjectInstance):ObjectName = that.getObjectName
  implicit def objectName2RichMBean(that:ObjectName)(implicit ajmx:JMX) = RichMBean(that)
  implicit def string2RichMBean(that:String)(implicit ajmx:JMX):RichMBean = RichMBean(that)
  implicit def toMarshalledObject[A](that:A):MarshalledObject[A] = new MarshalledObject[A](that)

  implicit def compositeDataAsScalaWrapper(cd : CompositeData): CompositeDataWrapper = CompositeDataWrapper(cd)
  implicit def compositeDataSupportAsScalaWrapper(cd : CompositeDataSupport): CompositeDataSupportWrapper = CompositeDataSupportWrapper(cd)
  implicit def tubularDataAsScalaWrapper(td:TabularData):TabularDataWrapper = TabularDataWrapper(td)
  implicit def tubularDataSupportAsScalaWrapper(td:TabularDataSupport):TabularDataSupportWrapper = TabularDataSupportWrapper(td)

}

object JMXImplicits extends JMXImplicits {
}



object RichMBean extends JMXImplicits {
  def apply(name:String)(implicit jmx:JMX) = new RichMBean(jmx, name)
  def apply(objectName:ObjectName)(implicit jmx:JMX) = new RichMBean(jmx, objectName)
}


 
class RichMBean(jmx:JMX, val objectName:ObjectName) extends JMXImplicits {
  
  implicit private val objectInstance = jmx.rmiConnection.getObjectInstance(objectName, null) match {
    case null  => throw new RuntimeException("Couldn't find MBean")
    case mbean => mbean
  }
  val name = objectName.toString()
  //val keys = mapAsScalaMap(objectName.getKeyPropertyList())  // Not available in scala 2.8.x
  val keys = (objectName.getKeyPropertyList().entrySet().toList map {case entry => entry.getKey() -> entry.getValue()}).toMap
  
  def mbean = this
  def getOption[A](name:String):Option[A] = {
    val mbeanInfo = jmx.rmiConnection.getMBeanInfo(objectName, null)
    val attrsInfo:List[javax.management.MBeanAttributeInfo] = mbeanInfo.getAttributes().toList
    (attrsInfo find {attrInfo => attrInfo.getName == name}) match {
      case None => None
      case Some(attrInfo) => get[A](name) match {
          case x if (x==null) => None
          case x => Some(x)
      }
    }
  }
  def get[A](attrname:String) = jmx.rmiConnection.getAttribute(objectName,attrname,null).asInstanceOf[A]
  def set[A](attrname:String, value:A) {
    val attribute = new javax.management.Attribute(attrname, value)
    jmx.rmiConnection.setAttribute(objectName, attribute, null)
  }
  def get[A](attrnames:String*):List[A] = attrnames.toList map { get[A](_)}
  def call[A](operation:String, args:Any*) : Option[A] = {
    val signature=args map { // http://home.pacifier.com/~mmead/jni/cs510ajp/index.html  /   http://en.wikipedia.org/wiki/Java_Native_Interface#Mapping_types
      case _:Int     => "int"
      case _:Short   => "short"
      case _:Long    => "long"
      case _:Float   => "float"
      case _:Double  => "double"
      case _:Boolean => "boolean"
      case _:Byte    => "byte"
      case _:Char    => "char"
      case _:String  => "String"
      case _:Array[Int]     => "[I"
      case _:Array[Short]   => "[S"
      case _:Array[Long]    => "[J"
      case _:Array[Float]   => "[F"
      case _:Array[Double]  => "[D"
      case _:Array[Boolean] => "[Z"
      case _:Array[Byte]    => "[B"
      case _:Array[Char]    => "[C"
      case _:Array[String]  => "[Ljava/lang/String"
    }
    val args2give=args.toArray
    jmx.rmiConnection.invoke(objectName, operation, args2give, signature.toArray, null) match {
      case null => None
      case result => Some(result.asInstanceOf[A])
    }
  }
  def browse() = {
    val mbeanInfo = jmx.rmiConnection.getMBeanInfo(objectName, null)
    val attrsInfo:List[javax.management.MBeanAttributeInfo] = mbeanInfo.getAttributes().toList
    attrsInfo map { _.getName }
  }
  //def apply[A](attrname:String) = jmx.rmiConnection.getAttribute(objectName,attrname,null).asInstanceOf[A]
  
}


case class Credentials(username:String="", password:String="")

case class JMXOptions(
  host: String = "localhost",
  port: Int = 22,
  username: Option[String] = None,
  password: Option[String] = None,
  connectTimeout: Int = 30000,
  retryCount: Int = 5,
  retryDelay: Int = 2000
  ) {
  val credentials:Option[Credentials] = (username,password) match {
    case (None,None) => None
    case (None,_) => None
    case (Some(u), Some(p)) => Some(Credentials(u,p))
  }
}



object JMX extends JMXImplicits with JMXAutoClose {
  def register(ob:Object, obname:ObjectName) = ManagementFactory.getPlatformMBeanServer.registerMBean(ob, obname)
  
  def connect[R](
      host:String,
      port:Int,
      username:Option[String]=None,
      password:Option[String]=None) (block: JMX => R):R = {
    connect[R](JMXOptions(host=host, port=port, username=username, password=password))(block)
  }
  def connect[R](options:JMXOptions) (block: JMX => R) = {
    usingJMX(new JMX(options)) { block(_)}
  }
  def connect[R](someOptions:Option[JMXOptions]) (block: JMX => R) = someOptions map { connect[R](_)(block)} 
    
  def apply(options:JMXOptions) = new JMX(options)
}





class JMX(val options:JMXOptions) extends JMXImplicits {
    
  def jmx=this
  val rmiConnection = findMBeanServers(options.host, options.port, options.credentials).head
  private def findMBeanServers(host:String, port:Int, credentials:Option[Credentials]) :List[RMIConnection] = {
    try {
      val lr = java.rmi.registry.LocateRegistry.getRegistry(host, port)
      var entries = lr.list.toList.filterNot { _ contains "TMFactory" } //=> This MBean causes registry data corruption if lookuped then all others call are in exception
      entries map { entry =>
        try {
          val credentialsArray:Array[String] = credentials match {
            case None => Array()
            case Some(Credentials(username, password)) => Array(username, password)
          }
          lr.lookup(entry) match {
            case stub:RMIServerImpl_Stub => Some(stub.newClient(credentialsArray))
            case rsi:RMIServer => Some(rsi.newClient(credentialsArray))
            case r@_ => println("Do not know how to manage %s class %s".format(entry, r.getClass)); None
          }
        } catch {
          case e:UnmarshalException => /*println("%s - Missing = %s".format(entry, searchMissingClass(e))); */ None
        }
      } flatten
    }
  }
  def close() = {
    try {
      rmiConnection.close
    } catch {
      case e => println("Exception while closing rmi connection, let's ignore and continue...")
    }
  }
  def searchMissingClass(e:Throwable):String = {
    e match {
      case null => "?"
      case ex:ClassNotFoundException => ex.getMessage().trim().split("\\s+")(0)
      case _ => searchMissingClass(e.getCause)
    }
  }
  def apply(name:String) = get(name)
  def get(name:String) = new RichMBean(this, name)
  def domains = rmiConnection.getDomains(null).toList
  def getOption(name:String) = if (rmiConnection.queryNames(name, null, null).size>0) Some(get(name)) else None
  //
  def query(query:String) = {
    rmiConnection.queryNames(null,query,null).toList map {RichMBean(_)(this)}
  }
  def browse() = {
    rmiConnection.queryNames(null,null,null).toList
  }
  def os()        = getOption("java.lang:type=OperatingSystem")
  def runtime()   = getOption("java.lang:type=Runtime")
  def threading() = getOption("java.lang:type=Threading")
  def memory()    = getOption("java.lang:type=Memory")
  /*
  def systemProperties():Map[String,String] = {
    runtime().map { r => r.get[TabularDataSupport]("SystemProperties") collect {
    	case (k:java.util.List[String],v:CompositeDataSupport) 
    		if (v.values.size==2) => k.head.toString -> v.values.tail.head.toString 
      }
    }  getOrElse Map.empty
  }
  */
  def systemProperties():Map[String,String] = {
    runtime() match {
      case Some(runtime) =>
        runtime.get[TabularDataSupport]("SystemProperties") flatMap {
          case (k:java.util.List[String],v:CompositeDataSupport) if (v.values.size==2) =>
            k.head.toString -> v.values.tail.head.toString :: Nil   // values first item is the system property name !
          case _ => Nil
        } toMap
      case None => Map.empty
    }
  }

}

